import numpy as np
import torch

from main import LstmRNN


def predict(test_x):
    lstm_model = LstmRNN(1, 16, output_size=1, num_layers=1)
    lstm_model.load_state_dict(torch.load('model_params.pkl'))  # load model parameters from files
    lstm_model = lstm_model.eval()  # switch to testing model

    # prediction on test dataset
    test_x_tensor = test_x.reshape(-1, 1,
                                   1)  # set batch size to 5, the same value with the training set
    test_x_tensor = torch.from_numpy(test_x_tensor)

    predictive_y_for_testing = lstm_model(test_x_tensor)
    predictive_y_for_testing = predictive_y_for_testing.view(-1, 1).data.numpy()
    print(predictive_y_for_testing)


if __name__ == '__main__':
    test_x = [6482]
    test_x = np.array(test_x).reshape(-1,1).astype('float32')
    predict(test_x)